#coding=utf8
# Copyright (c) 2016 Tinydot. inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import numpy as np
import edt
from scipy.spatial import KDTree
import os
from PIL import Image
from tbox.uniquecolors import ins_numpy_to_dis_image, K_RED, K_GRE, K_BLK, RED, GRE, BLK
from PIL import ImageDraw


def create_pos_map(w, h, norm=False):
    pos = np.where(np.zeros((w, h)) == 0)
    if norm:
        pos = (pos[0]/pos[0].astype(np.float).max(),
               pos[1]/pos[1].astype(np.float).max(),
               )
    return np.dstack(pos).reshape((w, h, 2))


def create_stack_pos_map(w, h, c, norm=False):
    pos = create_pos_map(w, h, norm)
    return np.stack([pos for _ in range(c)])


def find_edt_center(binary_image):
    assert len(binary_image.shape) == 2, \
        "input shape size must be (w, h), input: %s" % str(binary_image.shape)
    ins_pos = np.where(binary_image > 0)
    if ins_pos[0].shape[0] == 0:
        return None, None
    dt_dist = edt.edt(binary_image, black_border=True)
    peak_pos = np.dstack(np.where(dt_dist == dt_dist.max()))[0]
    tree = KDTree(peak_pos)
    mean_ctr = np.dstack(ins_pos).mean(-2)
    target = tree.data[tree.query(mean_ctr)[1]][0]
    return [int(_) for _ in target], round(float(dt_dist[tuple(target)]), 6)


def instance_edt_center_offset(instance_image, return_offset=True, min_count=-1):
    w, h = instance_image.shape
    ret_offset = None
    ret_static_pos = None
    if return_offset:
        ret_offset = np.zeros((w, h, 2)).astype(np.int)
        ret_static_pos = create_pos_map(*instance_image.shape, norm=False)
    meta = {}
    c_list = []
    if min_count > 0:
        u_list, c_list = np.unique(instance_image, return_counts=True)
    else:
        u_list = np.unique(instance_image)
    for i, u in enumerate(u_list):
        if u == 0:
            continue
        if min_count > 0:
            if c_list[i] < min_count:
                continue
        mask = instance_image == u
        assert isinstance(mask, np.ndarray)
        f_center, f_dist = find_edt_center(mask.astype(np.int))
        if return_offset:
            ret_offset[mask] = f_center - ret_static_pos[mask]
        meta[str(u)] = {"pos": f_center, "dis": f_dist}
    if return_offset:
        return ret_offset, meta
    return meta


def offset_to_image(offset_map):
    """
    000F0xxxxx0F0xxxxx, xxxxx = size_of(max_off)
    eg:
        +2048, -1234 =>  0020480101234
    """
    assert isinstance(offset_map, np.ndarray)
    w, h, c = offset_map.shape
    assert c == 2
    ret_left = np.zeros((w, h )).astype(np.int)
    ret_right = np.zeros((w, h)).astype(np.int)
    flag_pos = len("%s" % max([w, h])) + 1
    flag_val = int(10**flag_pos)

    def make_cp(z_target, val):
        z_target[val < 0] = flag_val
        return z_target + np.abs(val)

    ret_left = make_cp(ret_left, offset_map[:, :, 0])
    ret_right = make_cp(ret_right, offset_map[:, :, 1])
    return ret_left*flag_val*10 + ret_right


def image_to_offset(image_array):
    assert isinstance(image_array, np.ndarray)
    assert len(image_array.shape) == 2
    w, h = image_array.shape
    flag_pos = len("%s" % max([w, h])) + 1
    flag_val = int(10 ** flag_pos)
    assert np.max(image_array) >= flag_val

    def parse_value(value):
        ret_data = np.ones((w, h)).astype(np.int)
        ret_data[value >= flag_val] = -1
        return ret_data * (value % flag_val)

    val_right = parse_value(image_array % (flag_val*10))
    val_left = parse_value(image_array // (flag_val*10))
    return np.dstack([val_left, val_right])


def load_files_form_dir(target_dir, file_suffix):
    ret_list = []

    def __list_files(t_dir):
        for f in os.listdir(t_dir):
            fp = os.path.join(t_dir, f)
            if os.path.isdir(fp):
                __list_files(fp)
            else:
                if fp.endswith(file_suffix):
                    ret_list.append(fp)
    __list_files(target_dir)
    return ret_list


def load_edt_from_image(image_path):
    img = np.array(Image.open(image_path))
    return image_to_offset(img)


def save_edt_as_image(offset_map, image_path):
    img = offset_to_image(offset_map)
    img = Image.fromarray(img, mode="I")
    return img.save(image_path)


def draw_edt_meta(target_img, meta, no_arc=False, show_label=False):
    draw = ImageDraw.Draw(target_img)
    for k, v in meta.items():
        dis, (y, x) = v["dis"], v["pos"]
        if not no_arc:
            draw.arc((x - dis, y - dis, x + dis, y + dis), 0, 360, GRE)
        draw.point((x, y), RED)
        if show_label:
            draw.text((x, y), "%s" % k)
    return target_img


def mark_edt_as_image(raw_instance, heat_image, no_arc):
    offset_map, offset_meta = instance_edt_center_offset(raw_instance)
    if heat_image:
        image_data = np.linalg.norm(offset_map, axis=-1)
        image_max_value = np.max(image_data)
        image_data = 255*image_data/image_max_value
        mark_image = Image.fromarray(image_data.astype(np.uint8)).convert("RGB")
    else:
        mark_image = ins_numpy_to_dis_image(raw_instance, [K_RED, K_GRE, K_BLK])
    mark_image = draw_edt_meta(mark_image, offset_meta, no_arc)
    return mark_image, offset_meta


def comm_app_conv_edt_offset(arg_fc, img_process_fc, args, _debug = None):
    def default_debug(*a, **k):
        print(*a, **k)
    debug = _debug or default_debug
    import argparse
    import time
    import sys
    parser = argparse.ArgumentParser("edt offset map")
    parser.add_argument("instance_image",
                        help="instance_image or dirs", type=str)
    parser.add_argument("-s", "--show",
                        help="show marked image", action="store_true", default=False)
    parser.add_argument("--heat",
                        help="show marked image with heat model", action="store_true", default=False)
    parser.add_argument("--no_arc",
                        help="show marked image with no arc", action="store_true", default=False)
    parser.add_argument("-w", "--write",
                        help="write to file", action="store_true", default=False)
    parser = arg_fc(parser)
    ag = parser.parse_args(args)
    target = ag.instance_image
    targets_files = []
    if os.path.isfile(target):
        if ag.show:
            ins_array = img_process_fc(target)
            marked_image, meta = mark_edt_as_image(ins_array, ag.heat, ag.no_arc)
            dis_info = [v["dis"] for k, v in meta.items()]
            dis_info.sort()
            debug("dis_list[heat=%s]: %s" % (ag.heat, str(dis_info)))
            marked_image.show()
            if ag.write:
                f_name, f_suffix = os.path.splitext(target)
                marked_image.save("%s_mk%s" % (f_name, f_suffix))
            return
        targets_files = [target]
    elif os.path.isdir(target):
        targets_files = load_files_form_dir(target, ag.suffix)
    all_max_edt, all_min_edt = 0.0, sys.maxsize
    file_count = len(targets_files)
    debug("process %d files" % file_count)
    all_start = time.time()
    for i, f in enumerate(targets_files):
        debug("process[%d/%d] %s ..." % (i, file_count, f), end="\r")
        t_p = time.time()
        # process image
        image = img_process_fc(f)
        offset_map, offset_meta = instance_edt_center_offset(image)
        if ag.write:
            f_name, f_suffix = os.path.splitext(f)
            save_edt_as_image(offset_map, "%s_edt%s" % (f_name, f_suffix))
        dis_list = [d["dis"] for _, d in offset_meta.items()]
        min_edt, max_edt = -1, -1
        off_count = len(dis_list)
        if off_count:
            min_edt, max_edt = min(dis_list), max(dis_list)
            all_min_edt = min(all_min_edt, min_edt)
            all_max_edt = max(all_max_edt, max_edt)
        debug("process[%d/%d] %s complete[%d], time:\t%.2f, "
              "min_edt:\t%.2f/%.2f, max_edt:\t%.2f/%.2f" % (i, file_count, f, off_count,
                                                            time.time() - t_p,
                                                            min_edt, all_min_edt,
                                                            max_edt, all_max_edt))
    debug("ALL: time cost:\t%.2fs, all_min_edt:\t%.2f, all_max_edt:\t%.2f" % (
        time.time() - all_start, all_min_edt, all_max_edt
    ))


def app_cityscapes_edt_offset(args=None):

    def args_fc(parser):
        parser.add_argument("--suffix",
                            help="instance file suffix", type=str, default="_instanceIds.png")
        return parser

    def process_fc(image_file):
        image = np.array(Image.open(image_file))
        image[image % 1000 == 0] = 0
        image[image < 1000] = 0
        return image

    return comm_app_conv_edt_offset(args_fc, process_fc, args)
